use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::io;

async fn copy(s1: &mut TcpStream, s2: &mut TcpStream) {
    loop {
        let mut buf = [0; 4096];
        let n = match s1.read(&mut buf).await {
            // socket closed
            Ok(n) if n == 0 => return,
            Ok(n) => n,
            Err(e) => {
                eprintln!("failed to read from socket; err = {:?}", e);
                return;
            }
        };

        s2.write_all(&buf[0..n]).await.expect("cant open");
    }
}

pub async fn handle_connection(mut socket: TcpStream) {
    let mut buf = [0; 4096];
    let n = match socket.read(&mut buf).await {
        // socket closed
        Ok(n) if n == 0 => return,
        Ok(n) => n,
        Err(e) => {
            return;
        }
    };
    let mut server;
    //utf8编码 获取协议
    let req = String::from_utf8_lossy(&buf).to_string();
    let _data: Vec<&str> = req.split(" ").collect();
    //https
    if _data[0] == "CONNECT" {
        server = TcpStream::connect(_data[1]).await.expect("cant open");
        socket.write_all(b"HTTP/1.0 200 Connection Established\r\n\r\n").await;
    } else {
        //http
        let _data: Vec<&str> = req.split("\r\n").collect();
        let mut host = String::new();
        for line in _data {
            let v: Vec<&str> = line.split(": ").collect();
            if v[0].eq_ignore_ascii_case("host") {
                host = v[1].to_string();
                if !v[1].contains(":") {
                    host += ":80";
                }
                break;
            }
        }
        if host.is_empty() {
            return;
        }
        server = TcpStream::connect(host).await.expect("cant open");
        server.write_all(&buf[0..n]).await;
    }

    copy(&mut socket, &mut server);
    tokio::task::spawn_blocking(move || {
        copy(&mut server, &mut socket);
    }).await;
}